package twoD.IsoPfem;

import java.util.ArrayList;

import oneD.fem.algorithm.MyLinalg;
import inf.jlinalg.ArrayVector;
import inf.jlinalg.BLAM;
import inf.jlinalg.GeneralMatrixLSESolver;
import inf.jlinalg.ILSESolver;
import inf.jlinalg.IMatrix;
import inf.jlinalg.IVector;
import inf.jlinalg.QuadraticMatrixInfo;
import inf.jlinalg.SolveFailedException;

public abstract class Structure implements IStructure {

	private static final double TOL = 1e-6;

	public enum CType {
		PLANE_STRESS, PLANE_STRAIN, ISOMETRIC_3D
	}

	private Mesh mesh_;
	private int d_ = 2;
	protected int numberOfEquations_;
	private boolean isEnumerated = false;

	public Structure(Mesh M) {
		mesh_ = M;
	}

	/*
	 * TODO: add code for edges and faces
	 */

	@Override
	public Mesh getMesh() {
		return mesh_;
	}

	@Override
	public ArrayList<IElement> getElements() {
		ArrayList<IElement> E = new ArrayList<IElement>();
		for (int i = 0; i < getNumberOfElements(); i++)
			E.add(getElements(i));
		return E;
	}

	@Override
	public IElement getElements(IGeometry g) {
		for (int i = 0; i < getNumberOfElements(); i++)
			if (getElements(i).getGeometry() == g)
				return getElements(i);
		return null;
	}

	public void setCType(CType c) {
		for (int i = 0; i < getNumberOfElements(); i++)
			getElements(i).setCType(c);
	}

	/**
	 * 
	 * @return number of equations in Ku=F
	 */
	public int getNumberOfEquations() {
		return numberOfEquations_;
	}

	@Override
	public int enumerateDOFs() {
		if (!isEnumerated) {
			int start = 0;
			/*
			 * enumerate for nodes
			 */
			for (int i = 0; i < getMesh().getNumberOfNodes(); i++) {
				start = getMesh().getNodes(i).enumerateDOFs(start);
				getMesh().getNodes(i).setId(i);
			}
			/*
			 * enumerate for edges
			 */
			for (int i = 0; i < getMesh().getNumberOfEdges(); i++)
				start = getMesh().getEdges(i).enumerateDOFs(d_, start);
			/*
			 * enumerate for faces
			 */
			for (int i = 0; i < getMesh().getNumberOfFaces(); i++)
				start = getMesh().getFaces(i).enumerateDOFs(d_, start);

			numberOfEquations_ = start;

			/*
			 * enumerate for elements
			 */
			for (int i = 0; i < getNumberOfElements(); i++) {
				getElements(i).enumerateDOFs();
				getElements(i).setId(i);
			}

			isEnumerated = true;
			return start;
		} else
			return numberOfEquations_;
	}

	@Override
	public void assembleInternalForceVector(double[] riGlobal) {
		for (int i = 0; i < getNumberOfElements(); i++) {
			IVector rie = getElements(i).computeInternalForceVector();
			ArrayList<Integer> dofNumbers = getElements(i).getDOFNumbers();
			int l = dofNumbers.size();
			for (int p = 0; p < l; p++)
				if (dofNumbers.get(p) != -1)
					riGlobal[dofNumbers.get(p)] += rie.get(p);
		}
	}

	@Override
	public void assembleExternalForceVector(double[] rGlobal) {
		/*
		 * compute external forces per each element
		 */
		for (int i = 0; i < getNumberOfElements(); i++)
			getElements(i).computeExternalForceVector();
		/*
		 * assemble to global external forces
		 */
		for (int i = 0; i < getMesh().getNumberOfNodes(); i++) {
			int[] dofNumbers = getMesh().getNodes(i).getDofNumbers();
			for (int j = 0; j < d_; j++)
				if (dofNumbers[j] != -1) {
					rGlobal[dofNumbers[j]] += getMesh().getNodes(i)
							.getForceVector().get(j);
					rGlobal[dofNumbers[j]] += getMesh().getNodes(i)
							.getNodalForceVector().get(j);
				}
		}
		for (int i = 0; i < getMesh().getNumberOfEdges(); i++) {
			int[] dofNumbers = getMesh().getEdges(i).getDofNumbers();
			Force[] F = getMesh().getEdges(i).getNodalForceVector();
			if (F != null)
				for (int k = 0; k < F.length; k++)
					for (int j = 0; j < d_; j++)
						if (dofNumbers[k * d_ + j] != -1)
							rGlobal[dofNumbers[j]] += F[k].getComponent(j);
		}
		/*
		 * TODO: add code to assemble force for face
		 */
	}

	@Override
	public void assembleTangentStiffnessMatrix(IMatrix kGlobal, boolean full) {
		for (int i = 0; i < getNumberOfElements(); i++) {
			IMatrix ke = getElements(i).computeTangentStiffnessMatrix();
			// System.out.println("ke" + i + "=" + MatrixFormat.format(ke));
			ArrayList<Integer> dofNumbers = getElements(i).getDOFNumbers();
			int l = dofNumbers.size();
			for (int p = 0; p < l; p++)
				for (int q = full ? 0 : p; q < l; q++)
					if (dofNumbers.get(p) != -1 && dofNumbers.get(q) != -1)
						kGlobal.add(dofNumbers.get(p), dofNumbers.get(q),
								ke.get(p, q));
		}
	}

	/**
	 * solve (ri(u)-r=0) by iterate n times. It can be good for linear system
	 * where we just need to iterate one time
	 * 
	 * @param n
	 *            number of loops
	 */
	public void Iterate(int n) {
		int NEQ = enumerateDOFs();
		// solver initialization
		ILSESolver solver = new GeneralMatrixLSESolver(
				GeneralMatrixLSESolver.SOLVER_IMPL_JAVA);
		// ILSESolver solver = new SpoolesLSESolver();
		// using this for solver implementation in Java, it will be slow
		QuadraticMatrixInfo aInfo = solver.getAInfo();
		IMatrix kGlobal = solver.getA();
		double[] rGlobal = new double[NEQ];
		double[] riGlobal = new double[NEQ];
		double[] uGlobal = new double[NEQ];

		aInfo.setSize(NEQ);

		for (int i = 0; i < n; i++) {
			solver.initialize();

			assembleTangentStiffnessMatrix(kGlobal, true);
			assembleExternalForceVector(rGlobal);
			assembleInternalForceVector(riGlobal);

			// System.out.println("kGlobal = " + MatrixFormat.format(kGlobal));
			// System.out.println("rGlobal = " + ArrayFormat.format(rGlobal));

			MyLinalg.add(-1.0, riGlobal, rGlobal);

			// solve the system of equation
			try {
				solver.solve(rGlobal);
			} catch (SolveFailedException e) {
				System.out.println("Solver failed: " + e.getMessage());
			}

			// update the displacement
			MyLinalg.add(1.0, rGlobal, uGlobal);
			// System.out.println("Solution x:");
			// System.out.println(ArrayFormat.format(uGlobal));

			// pass the result back to nodes
			this.setDisplacement(uGlobal);

			// reset the immediate variables to proceed to next step
			BLAM.zero(kGlobal);
			MyLinalg.zero(rGlobal);
			MyLinalg.zero(riGlobal);
		}
	}

	/**
	 * solve (ri(u)-r=0) by iterate until converged
	 * 
	 * @return number of loop performed
	 */
	public int Iterate() {
		return Iterate(1.0);
	}

	/**
	 * solve (ri(u)-lambda*r=0) by iterate until converged
	 * 
	 * @return number of loop performed
	 */
	public int Iterate(double lambda) {
		int NEQ = enumerateDOFs();
		// solver initialization
		ILSESolver solver = new GeneralMatrixLSESolver(
				GeneralMatrixLSESolver.SOLVER_IMPL_JAVA);
		// ILSESolver solver = new SpoolesLSESolver();
		// using this for solver implementation in Java, it will be slow
		QuadraticMatrixInfo aInfo = solver.getAInfo();
		IMatrix kGlobal = solver.getA();
		double[] rGlobal = new double[NEQ];
		double[] riGlobal = new double[NEQ];
		double[] uGlobal = new double[NEQ];

		aInfo.setSize(NEQ);

		double eps = 1;
		int cnt = 0;
		while (eps > TOL) {
			solver.initialize();

			assembleTangentStiffnessMatrix(kGlobal, true);
			assembleExternalForceVector(rGlobal);
			if (lambda != 1.0)
				MyLinalg.scale(lambda, rGlobal);
			assembleInternalForceVector(riGlobal);

			// System.out.println("kGlobal = " + MatrixFormat.format(kGlobal));
			// System.out.println("rGlobal = " + ArrayFormat.format(rGlobal));

			MyLinalg.add(-1.0, riGlobal, rGlobal);

			// solve the system of equation
			try {
				solver.solve(rGlobal);
			} catch (SolveFailedException e) {
				System.out.println("Solver failed: " + e.getMessage());
			}

			// update the displacement
			MyLinalg.add(1.0, rGlobal, uGlobal);
			// System.out.println("Solution x:");
			// System.out.println(ArrayFormat.format(uGlobal));

			// pass the result back to nodes
			this.setDisplacement(uGlobal);

			// compute the relative error
			eps = MyLinalg.norm2(rGlobal) / MyLinalg.norm2(uGlobal);

			cnt++;

			// reset the immediate variables to proceed to next step
			BLAM.zero(kGlobal);
			MyLinalg.zero(rGlobal);
			MyLinalg.zero(riGlobal);
		}
		return cnt;
	}

	public void printStructure() {
		System.out.println("Listing structure\n");
		System.out.println("Nodes:");
		System.out.println("idx\t\tx1\t\tx2\t\tx3");
		for (int i = 0; i < getMesh().getNumberOfNodes(); i++) {
			System.out.print(i + "\t");
			getMesh().getNodes(i).printPosition();
			System.out.println();
		}
		System.out.println("Constraints:");
		System.out.println("node\tu1\tu2\tu3");
		for (int i = 0; i < getMesh().getNumberOfNodes(); i++) {
			System.out.print(i + "\t");
			getMesh().getNodes(i).printConstraint();
			System.out.println();
		}
		System.out.println("Forces:");
		System.out.println("node\t\tr1\t\tr2\t\tr3");
		for (int i = 0; i < getMesh().getNumberOfNodes(); i++) {
			if (getMesh().getNodes(i).getForce() != null) {
				System.out.print(i + "\t");
				getMesh().getNodes(i).getForce().print();
				System.out.println();
			}
		}
	}

	public void listingDisplacements() {
		System.out.println("Listing displacement\n");
		System.out.println("Nodes:");
		System.out.println("idx\t\tdx1\t\tdx2");
		for (int i = 0; i < getMesh().getNumberOfNodes(); i++) {
			System.out.print(i + "\t");
			getMesh().getNodes(i).printDisplacement();
			System.out.println();
		}
	}

	@Override
	public IVector getUGlobal() {
		IVector U = new ArrayVector(d_ * getMesh().getNumberOfNodes());
		int cnt = 0;
		for (int i = 0; i < getMesh().getNumberOfNodes(); i++) {
			int[] dofNumbers = getMesh().getNodes(i).getDofNumbers();
			for (int j = 0; j < d_; j++)
				if (dofNumbers[j] != -1)
					U.set(cnt++, getMesh().getNodes(i).getDisplacement().get(j));
				else
					U.set(cnt++, 0.0);
		}
		/*
		 * TODO: to be rewritten to account for mode on edges & faces
		 */
		return U;
	}

	@Override
	public double[] getU() {
		double[] U = new double[getNumberOfEquations()];
		for (int i = 0; i < getMesh().getNumberOfNodes(); i++) {
			Node n = getMesh().getNodes(i);
			int[] dofNumbers = n.getDofNumbers();
			for (int j = 0; j < dofNumbers.length; j++)
				if (dofNumbers[j] != -1)
					U[dofNumbers[j]] = n.getDisplacement().get(j);
		}
		/*
		 * TODO: to be rewritten to account for mode on edges & faces
		 */
		return U;
	}

	@Override
	public IVector getXGlobal() {
		int numNode = getMesh().getNumberOfNodes();
		int numEdge = getMesh().getNumberOfEdges();
		int edgeMode = getMesh().getEdges(0).getNumberOfModes();
		int numFace = getMesh().getNumberOfFaces();
		int faceMode = getMesh().getFaces(0).getNumberOfModes();
		IVector X = new ArrayVector(d_
				* (numNode + edgeMode * numEdge + faceMode * numFace));
		int cnt = 0;
		for (int i = 0; i < numNode; i++) {
			double[] xi = getMesh().getNodes(i).getPosition().toArray();
			for (int j = 0; j < d_; j++)
				X.set(cnt++, xi[j]);
		}
		/*
		 * TODO: rewrite code to read point from edges and faces
		 */
		// for (int i = 0; i < numEdge; i++)
		// for (int j = 0; j < edgeMode; j++) {
		// double[] xi = getMesh().getEdges(i).getPoint();
		// for (int k = 0; k < d_; j++)
		// X.set(cnt++, xi[k]);
		// }
		// for (int i = 0; i < numFace; i++)
		// for (int j = 0; j < faceMode; j++) {
		// double[] xi = getMesh().getFaces(i).getPoint();
		// for (int k = 0; k < d_; j++)
		// X.set(cnt++, xi[k]);
		// }

		return X;
	}

	@Override
	public void setDisplacement(double[] uGlobal) {
		this.setDisplacement(new ArrayVector(uGlobal));
	}

	public void setDisplacement(IVector uGlobal) {
		/*
		 * uGlobal is displacement vector after constrainted
		 */
		for (int i = 0; i < getMesh().getNumberOfNodes(); i++) {
			double[] u = new double[d_];
			int[] dofNumbers = getMesh().getNodes(i).getDofNumbers();
			for (int j = 0; j < d_; j++)
				if (dofNumbers[j] == -1)
					u[j] = 0;
				else
					u[j] = uGlobal.get(dofNumbers[j]);
			getMesh().getNodes(i).setDisplacement(u);
		}
		for (int i = 0; i < getMesh().getNumberOfEdges(); i++) {
			double[] u = new double[d_
					* getMesh().getEdges(i).getNumberOfModes()];
			int[] dofNumbers = getMesh().getEdges(i).getDofNumbers();
			for (int j = 0; j < u.length; j++)
				if (dofNumbers[j] == -1)
					u[j] = 0;
				else
					u[j] = uGlobal.get(dofNumbers[j]);
			getMesh().getEdges(i).setDisplacement(u);
		}
		for (int i = 0; i < getMesh().getNumberOfFaces(); i++) {
			double[] u = new double[d_
					* getMesh().getFaces(i).getNumberOfModes()];
			int[] dofNumbers = getMesh().getFaces(i).getDofNumbers();
			for (int j = 0; j < u.length; j++)
				if (dofNumbers[j] == -1)
					u[j] = 0;
				else
					u[j] = uGlobal.get(dofNumbers[j]);
			getMesh().getFaces(i).setDisplacement(u);
		}
	}
}
